import cv2 as cv
import numpy as np


class RoadDetect:
    """车道线检测"""
    def __init__(self):
        self.video = None
        self.video_img = None
        self.fps = None
        self.height = None
        self.width = None
        self.init_video()

    def init_img(self):
        pass

    def init_video(self):
        video = cv.VideoCapture("img/road.mp4")
        if video.isOpened():
            self.video = video

    def process(self):
        if self.video is None:
            return

        # 获取图片的信息:帧率
        self.fps = self.video.get(cv.CAP_PROP_FPS)
        # 获取每帧宽度
        self.width = int(self.video.get(cv.CAP_PROP_FRAME_WIDTH))
        # 获取每帧的高度
        self.height = int(self.video.get(cv.CAP_PROP_FRAME_HEIGHT))

        while self.video.isOpened():
            res, self.video_img = self.video.read()

            # 发现边沿以及处理
            canny_img = self.find_canny(self.video_img)
            # 霍夫直线变换
            lines = self.hough_lines(canny_img)
            # 直线绘制
            # img_copy = img.copy()
            # for line in lines:
            #     x1, y1, x2, y2 = line[0]
            #     cv.line(img_copy, (x1, y1), (x2, y2), (0, 0, 255), 2)
            # cv.imshow("hough_img", img_copy)
            self.draw_lines(lines)
            cv.imshow("self.video_img", self.video_img)

            if res:
                cv.waitKey(int(1000 / self.fps))

    def find_canny(self, img):
        # 使用canny进行边沿检测
        canny_img = cv.Canny(img, 50, 100)
        # 保留感兴趣的区域
        mask_img = np.zeros_like(canny_img)
        polygons = np.array([[(200, self.height), (1100, self.height), (550, 250)]])
        # 生成一个多边形的白色块，覆盖我们想要的部分
        cv.fillPoly(mask_img, polygons, 255)
        # 在原图中选取mask中白块覆盖的部分,把不是白块的地方都涂黑
        canny_img[mask_img == 0] = 0
        # cv.imshow("canny_img", canny_img)
        return canny_img

    def hough_lines(self, img):
        """霍夫直线变换"""
        #  线段以像素为单位的距离精度，double类型的，推荐用1.0
        rho = 1
        # 线段以弧度为单位的角度精度，推荐用numpy.pi/180
        theta = np.pi / 180
        # 累加平面的阈值参数，int类型，超过设定阈值才被检测出线段，值越大，基本上意味着检出的线段越长，检出的线段个数越少。
        threshold = 10
        # 线段以像素为单位的最小长度
        min_line_length = 25
        # 同一方向上两条线段判定为一条线段的最大允许间隔（断裂），超过了设定值，则把两条线段当成一条线段，值越大，允许线段上的断裂越大，越有可能检出潜在的直线段
        max_line_gap = 3
        lines = cv.HoughLinesP(img, rho, theta, threshold, minLineLength=min_line_length, maxLineGap=max_line_gap)
        return lines

    def draw_lines(self, lines):
        """绘制车道两边的直线"""
        # 记录霍夫变换识别出的左右两侧直线的斜率和截距
        left_lines = []
        right_lines = []
        for line in lines:
            x1, y1, x2, y2 = line[0]
            # 两点确定一条直线
            params = np.polyfit((x1, x2), (y1, y2), 1)
            # 斜率
            slope = params[0]
            intercept = params[1]
            # 由于坐标系是向着右下角方式建立的
            # 斜率>0是右边直线
            if slope > 0:
                right_lines.append((slope, intercept))
            else:
                left_lines.append((slope, intercept))

        if not left_lines or not right_lines:
            return
        # 计算斜率和截距的平均值，以此来确定最终的直线
        left_lines = np.array(left_lines)
        right_lines = np.array(right_lines)
        slope_average_l, intercept_average_l = np.average(left_lines, axis=0)
        slope_average_r, intercept_average_r = np.average(right_lines, axis=0)

        # 根据平均斜率和截距绘制直线 y = slope * x + intercept
        # 求左侧直线的两点
        y1_l = self.height
        x1_l = int((y1_l - intercept_average_l) / slope_average_l)
        y2_l = int(self.height / 2)
        x2_l = int((y2_l - intercept_average_l) / slope_average_l)

        # 求右侧直线的两点
        y1_r = self.height
        x1_r = int((y1_r - intercept_average_r) / slope_average_r)
        y2_r = int(self.height / 2)
        x2_r = int((y2_r - intercept_average_r) / slope_average_r)

        # 绘制左侧直线
        cv.line(self.video_img, (x1_l, y1_l), (x2_l, y2_l), (0, 0, 255), 2)
        # 绘制右侧直线
        cv.line(self.video_img, (x1_r, y1_r), (x2_r, y2_r), (0, 0, 255), 2)


if __name__ == '__main__':
    RoadDetect().process()

    # 当程序推出时，释放所有窗口资源
    cv.destroyAllWindows()